from flask import session, redirect, url_for, g
from functools import wraps
import config
from .models import *


def login_need(func):
    @wraps(func)
    def need(*args, **kwargs):
        if config.CMS_USER_ID in session:
            user_id = session.get(config.CMS_USER_ID)
            user = CMSUser.query.filter_by(token=user_id).first()
            if user:
                g.cms_user = user
                perms = CMSRole.query.filter_by(cms_id=user.id).all()
                menu_name = ''
                for i in perms:
                    menu_name += i.menu_name
                    menu_name += ','
                g.cms_menu = menu_name
                # print(g.cms_menu)
                return func(*args, **kwargs)
            else:
                return redirect(url_for('cms.login'))
        else:
            return redirect(url_for('cms.login'))
    return need
